import pandas as pd
import numpy as np
import gerber

data1 = pd.read_csv(r'data\TRD_Dalyr.csv')
data2 = pd.read_csv(r'data\TRD_Dalyr1.csv')
print(data1.shape)
print(data2.shape)
data = pd.concat([data1, data2])  # 纵向合并两个数据集（默认axis=0），创建新DataFrame
print(data)
data['LogDretwd'] = np.log(data['Dretwd'] + 1)  # 对'Dretwd'列进行对数转换（加1避免对0或负值取对数）
print(data)

stock_codes = data['Stkcd'].unique()  # 使用.unique()方法提取'Stkcd'列中的唯一值，确保每个股票只出现一次
print(stock_codes)

# 创建空DataFrame，索引为唯一的交易日期（去重后的日期列表）
# 使用.unique()确保每个交易日只出现一次，作为行索引
# 该DataFrame将用于存储各股票的对数收益率时间序列
df = pd.DataFrame(index=data['Trddt'].unique())
for stkcd in stock_codes:
    temp = data[data['Stkcd'] == stkcd]  # 筛选出当前股票的所有交易记录（布尔索引过滤）
    if temp.shape[0] == df.shape[0]:
        # 将当前股票的对数收益率列添加到df中
        # 使用.values将pandas Series转换为numpy数组，提升赋值效率
        df[stkcd] = temp['LogDretwd'].values
print(df)

df = df.iloc[:, :500]  # 选取前500只股票，减少运行时间

# 计算各股票的平均收益率（按列计算均值）
# 结果为Series，索引为股票代码，值为对应平均收益率
mean_returns = df.mean()
print(mean_returns)
mean_returns.to_csv('./temp_files/mean_returns.csv')
# 计算各股票收益率的标准差（按列计算标准差）
# 标准差衡量收益率的波动性（风险）
std_returns = df.std()
print(std_returns)
std_returns.to_csv('./temp_files/std_returns.csv')
# 计算收益率的协方差矩阵
# 协方差矩阵反映股票间的收益联动性：
# - 对角线元素为各股票收益率的方差
# - 非对角线元素为两两股票的协方差
# 该矩阵是投资组合优化的核心输入参数
cov_matrix_HC = df.cov()  # 历史协方差计算
print("cov_matrix_HC")
print(cov_matrix_HC)
cov_matrix_HC.to_csv('cov_matrix_HC.csv')

cov_matrix_GB = gerber.get_gerber_matrix(df, 0.5)
print("cov_matrix_GB")
print(cov_matrix_GB)
cov_matrix_GB.to_csv('cov_matrix_GB.csv')

cov_matrix_MGB = gerber.get_gerber_matrix_modified(df, 0.5, 0.1, -0.05)  # lambda1:NN,lambda2:NU/ND/DN/UN
print("cov_matrix_MGB")
print(cov_matrix_MGB)
cov_matrix_MGB.to_csv('cov_matrix_MGB.csv')


# 定义风险计算函数（投资组合标准差）
# 参数：
#   x - 权重向量（各股票配置比例）
#   cov_matrix - 收益率协方差矩阵
# 公式：σ_p = √(xᵀΣx)，其中Σ为协方差矩阵
def getRisk(x, cov_matrix):
    return np.sqrt(np.dot(x, np.dot(cov_matrix, x)))


# 定义收益计算函数（投资组合预期收益）
# 参数：
#   x - 权重向量
#   mean_returns - 各股票平均收益率
# 公式：E(R_p) = xᵀμ，其中μ为平均收益率向量
def getReturn(x, mean_returns):
    return np.dot(x, mean_returns)


# 获取股票数量（即DataFrame的列数）
stock_num = df.shape[1]
print(stock_num)

# 创建等权重投资组合（均匀分配资金）
# 初始权重向量：每个股票分配相同比例（1/N）
# 这是最简单的配置策略，后续可作为优化起点
initiate_x = [1 / stock_num] * stock_num

from scipy import optimize

# 初始化有效前沿存储列表
# 每个元素格式为：(目标收益率, 对应最小风险, 权重向量)
frontier_HC = list()
frontier_GB = list()
frontier_MGB = list()

# 生成20个等间距的目标收益率（从最低到最高平均收益率）
# 使用np.linspace确保均匀覆盖整个收益范围
target_return_range = np.linspace(mean_returns.min(), mean_returns.max(), 20)
for target_return in target_return_range:
    print(f"正在计算目标收益率: {target_return:.4f}")

    # 设置权重约束条件（不允许卖空，即权重在0-1之间）
    bounds = tuple((0, 1) for _ in range(stock_num))

    # 定义优化约束条件：
    # 1. 收益率等式约束：投资组合收益必须等于目标值
    # 2. 权重总和等式约束：权重之和必须等于1
    constraints = (
        {'type': 'eq', 'fun': lambda x: target_return - getReturn(x, mean_returns)},  # 收益约束
        {'type': 'eq', 'fun': lambda x: 1 - sum(x)}  # 权重总和约束
    )
    # 使用历史数据法计算协方差矩阵
    # 执行风险最小化优化（使用序列最小二乘规划算法SLSQP）
    # 参数说明：
    # - objective: 目标函数（最小化风险）
    # - x0: 初始权重向量（等权重）
    # - args: 传递给目标函数的额外参数（协方差矩阵）
    # - method: 优化算法选择
    # - bounds: 变量边界约束
    # - constraints: 等式约束条件
    minimize_risk_HC = optimize.minimize(
        getRisk,
        initiate_x,
        args=cov_matrix_HC,
        method='SLSQP',
        bounds=bounds,
        constraints=constraints
    )

    # 更新初始权重为当前优化结果（用于下次迭代的暖启动）
    initiate_x = minimize_risk_HC.x

    # 记录有效前沿点（目标收益、最小风险、权重向量）
    frontier_HC.append([target_return, minimize_risk_HC.fun] + minimize_risk_HC.x.tolist())

pd_frontier_HC = pd.DataFrame(frontier_HC, index=target_return_range,
                              columns=['target_return', 'min_risk'] + stock_codes[:50].tolist())
print(pd_frontier_HC.shape)
print(pd_frontier_HC)
pd_frontier_HC.to_csv('./frontier/frontier_HC.csv')

for target_return in target_return_range:
    print(f"正在计算目标收益率(GB): {target_return:.4f}")

    # 设置权重约束条件（不允许卖空，即权重在0-1之间）
    bounds = tuple((0, 1) for _ in range(stock_num))

    # 定义优化约束条件：
    # 1. 收益率等式约束：投资组合收益必须等于目标值
    # 2. 权重总和等式约束：权重之和必须等于1
    constraints = (
        {'type': 'eq', 'fun': lambda x: target_return - getReturn(x, mean_returns)},  # 收益约束
        {'type': 'eq', 'fun': lambda x: 1 - sum(x)}  # 权重总和约束
    )
    # 使用GB计算协方差矩阵
    # 执行风险最小化优化（使用序列最小二乘规划算法SLSQP）
    # 参数说明：
    # - objective: 目标函数（最小化风险）
    # - x0: 初始权重向量（等权重）
    # - args: 传递给目标函数的额外参数（协方差矩阵）
    # - method: 优化算法选择
    # - bounds: 变量边界约束
    # - constraints: 等式约束条件
    minimize_risk_GB = optimize.minimize(
        getRisk,
        initiate_x,
        args=cov_matrix_GB,
        method='SLSQP',
        bounds=bounds,
        constraints=constraints
    )

    # 更新初始权重为当前优化结果（用于下次迭代的暖启动）
    initiate_x = minimize_risk_GB.x

    # 记录有效前沿点（目标收益、最小风险、权重向量）
    frontier_GB.append([target_return, minimize_risk_GB.fun] + minimize_risk_GB.x.tolist())

pd_frontier_GB = pd.DataFrame(frontier_GB, index=target_return_range,
                              columns=['target_return', 'min_risk'] + stock_codes[:500].tolist())
print(pd_frontier_GB.shape)
print(pd_frontier_GB)
pd_frontier_HC.to_csv('./frontier/frontier_GB.csv')

for target_return in target_return_range:
    print(f"正在计算目标收益率(MGB): {target_return:.4f}")

    # 设置权重约束条件（不允许卖空，即权重在0-1之间）
    bounds = tuple((0, 1) for _ in range(stock_num))

    # 定义优化约束条件：
    # 1. 收益率等式约束：投资组合收益必须等于目标值
    # 2. 权重总和等式约束：权重之和必须等于1
    constraints = (
        {'type': 'eq', 'fun': lambda x: target_return - getReturn(x, mean_returns)},  # 收益约束
        {'type': 'eq', 'fun': lambda x: 1 - sum(x)}  # 权重总和约束
    )
    # 使用MGB计算协方差矩阵
    # 执行风险最小化优化（使用序列最小二乘规划算法SLSQP）
    # 参数说明：
    # - objective: 目标函数（最小化风险）
    # - x0: 初始权重向量（等权重）
    # - args: 传递给目标函数的额外参数（协方差矩阵）
    # - method: 优化算法选择
    # - bounds: 变量边界约束
    # - constraints: 等式约束条件
    minimize_risk_MGB = optimize.minimize(
        getRisk,
        initiate_x,
        args=cov_matrix_MGB,
        method='SLSQP',
        bounds=bounds,
        constraints=constraints
    )

    # 更新初始权重为当前优化结果（用于下次迭代的暖启动）
    initiate_x = minimize_risk_MGB.x

    # 记录有效前沿点（目标收益、最小风险、权重向量）
    frontier_MGB.append([target_return, minimize_risk_MGB.fun] + minimize_risk_MGB.x.tolist())

pd_frontier_MGB = pd.DataFrame(frontier_MGB, index=target_return_range,
                               columns=['target_return', 'min_risk'] + stock_codes[:500].tolist())
print(pd_frontier_MGB.shape)
print(pd_frontier_MGB)
pd_frontier_MGB.to_csv('./frontier/frontier_MGB.csv')
